#include "parser/lexer.hpp"
#include "parser/sourceRange.hpp"
#include "util/arraylist.hpp"
#include "util/bufferedInputStream.hpp"
#include <cstdio>
#include <cstdlib>
#include <string>

Token::Token(TokenType tt, const char* v, int len, SourceRange range) {
    _range = range;
    _tt = tt;
    if (len != 0) {
        _value = new char[len+1];;
        _value[len] = 0;
        strncpy(_value, v, len);
        _length = len;
    }
    else {
        _value = NULL;
        _length = 0;
    }
}

Token::~Token() {
    delete[] _value;
    _value = NULL;
    _length = 0;
}

Lexer::Lexer(BufferedInputStream* stream) {
    _stream = stream;
    _state = INIT;
    _index = 0;
    resetLocation();
}

Lexer::~Lexer() {
    _stream = NULL;
    _state = INIT;
    _index = 0;
}

void Lexer::resetLocation() {
    _location = SourceLocation(1, 1);
}

char Lexer::filter_char(const char* v, int size) {
    enum {
        INIT,
        NORMAL,
        ESCAPE,
        ESCAPE_UNIQUE,
        ESCAPE_OCT,
        ESCAPE_HEX
    } state = INIT;
    int idx = 0;
    if (size > 4) {
        printf("Compilation Error: Illegal Char '%s'\n", v);
        exit(-1);
    }
    int oct_count = 0;
    int hex_count = 0;
    char res = '\0';
    while (idx < size) {
        char c = v[idx];
        if (state == INIT) {
            if (c == '\\') {
                state = ESCAPE;
            } else {
                state = NORMAL;
                res = c;
            }
            idx++;
        } else if (state == NORMAL) {
            printf("Compilation Error: Illegal Char '%s'\n", v);
            exit(-1);
        } else if (state == ESCAPE) {
            switch (c) {
                case 'a':
                    res = '\a';
                    state = ESCAPE_UNIQUE;
                    break;
                case 'b':
                    res = '\b';
                    state = ESCAPE_UNIQUE;
                    break;
                case 'f':
                    res = '\f';
                    state = ESCAPE_UNIQUE;
                    break;
                case 'n':
                    res = '\n';
                    state = ESCAPE_UNIQUE;
                    break;
                case 'r':
                    res = '\r';
                    state = ESCAPE_UNIQUE;
                    break;
                case 't':
                    res = '\t';
                    state = ESCAPE_UNIQUE;
                    break;
                case 'v':
                    res = '\v';
                    state = ESCAPE_UNIQUE;
                    break;
                case '\\':
                    res = '\\';
                    state = ESCAPE_UNIQUE;
                    break;
                case '?':
                    res = '\?';
                    state = ESCAPE_UNIQUE;
                    break;
                case '\'':
                    res = '\'';
                    state = ESCAPE_UNIQUE;
                    break;
                case '\"':
                    res = '\"';
                    state = ESCAPE_UNIQUE;
                    break;
                // '0-7'
                // 'x'
                case '7':
                case '6':
                case '5':
                case '4':
                case '3':
                case '2':
                case '1':
                case '0':
                    oct_count = c - '0';
                    state = ESCAPE_OCT;
                    break;
                case 'x':
                    state = ESCAPE_HEX;
                    break;
                default:
                    printf("Compilation Error: Illegal Char '%s'\n", v);
                    exit(-1);
            }
            idx ++;
        } else if (state == ESCAPE_UNIQUE) {
            printf("Compilation Error: Illegal Char '%s'\n", v);
            exit(-1);
        } else if (state == ESCAPE_OCT) {
            if (c < '0' || c > '7') {
                printf("Compilation Error: Illegal Char '%s'\n", v);
                exit(-1);
            }
            oct_count = oct_count * 8 + (c - '0');
            idx++;
        } else if (state == ESCAPE_HEX) {
            if ((c >= '0' && c <= '9')
                || (c >= 'a' && c <= 'f')
                || (c >= 'A' && c <= 'F')) {
                hex_count = hex_count * 16 + std::stoi(std::string(1, c), nullptr, 16);
                idx++;
            } else {
                printf("Compilation Error: Illegal Char '%s'\n", v);
                exit(-1);
            }
        }
    }
    switch (state) {
        case INIT:
        case ESCAPE:
            printf("Error state: %d\n", state);
            exit(-1);
        case NORMAL:
        case ESCAPE_UNIQUE:
            return res;
        case ESCAPE_OCT:
            return static_cast<char>(oct_count);
        case ESCAPE_HEX:
            return static_cast<char>(hex_count);
    }
    return '\0';
}

Token* Lexer::filter_keyword(Token* t) {
    if (t->cmp("if") == 0) {
        t->_tt = TokenType::IF;
    }
    else if (t->cmp("else") == 0) {
        t->_tt = TokenType::ELSE;
    }
    else if (t->cmp("var") == 0) {
        t->_tt = TokenType::VAR;
    }
    else if (t->cmp("Int") == 0) {
        t->_tt = TokenType::TYPE_INT;
    }
    else if (t->cmp("Bool") == 0) {
        t->_tt = TokenType::TYPE_BOOL;
    }
    else if (t->cmp("String") == 0) {
        t->_tt = TokenType::TYPE_STRING;
    }
    else if (t->cmp("Double") == 0) {
        t->_tt = TokenType::TYPE_DOUBLE;
    }
    else if (t->cmp("Char") == 0) {
        t->_tt = TokenType::TYPE_CHAR;
    }
    else if (t->cmp("type") == 0) {
        t->_tt = TokenType::TYPE;
    }
    else if (t->cmp("Any") == 0) {
        t->_tt = TokenType::TYPE_ANY;
    }
    else if (t->cmp("True") == 0) {
        t->_tt = TokenType::BOOL_TRUE;
    }
    else if (t->cmp("False") == 0) {
        t->_tt = TokenType::BOOL_FALSE;
    }
    return t;
}

void Lexer::lexSkipCharacter() {
    while (true) {
        char skip = _stream->peek();
        if (skip == ' ' || skip == '\t' || skip == '\n' || skip == '\r') {
            _stream->advance();
            _location.moveToNext(skip);
            continue;
        }
        break;
    }
}

Token* Lexer::next() {
    lexSkipCharacter();
    SourceLocation begin = SourceLocation(_location);
    SourceLocation end = SourceLocation(_location);
    char c = _stream->peek();
    if (c == EOF) {
        return new Token(TokenType::ENDMARKER, "", 0, SourceRange(begin, end));
    }

    ArrayList<char> list_token(4);

    while (true) {
        c = _stream->peek();
        if (_state == INIT) {
            _stream->advance();
            if (c >= '0' && c <= '9') {
                _state = INTEGER;
                list_token.add(c);
                end = _location;
                _location.moveToNext(c);
            }
            else if ((c >= 'a' && c <= 'z')
                || (c >= 'A' && c <= 'Z')
                || c == '_') {
                _state = NAME;
                list_token.add(c);
                end = _location;
                _location.moveToNext(c);
            }
            else if (c == '=') {
                _state = OP_EQ;
                list_token.add(c);
                end = _location;
                _location.moveToNext(c);
            }
            else if (c == '$') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::LAMBDA, "$", 1, SourceRange(begin, end));
            }
            else if (c == ',') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::COMMA, ",", 1, SourceRange(begin, end));
            }
            else if (c == '!') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::LOGIC_NOT, "!", 1, SourceRange(begin, end));
            }
            else if (c == '(') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::LEFT_PAR, "(", 1, SourceRange(begin, end));
            }
            else if (c == ')') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::RIGHT_PAR, ")", 1, SourceRange(begin, end));
            }
            else if (c == '{') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::LEFT_BRACKET, "{", 1, SourceRange(begin, end));
            }
            else if (c == '}') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::RIGHT_BRACKET, "}", 1, SourceRange(begin, end));
            }
            else if (c == '+') {
                _state = OP_ADD;
                end = _location;
                _location.moveToNext(c);
                list_token.add(c);
            }
            else if (c == '-') {
                _state = OP_SUB;
                end = _location;
                _location.moveToNext(c);
                list_token.add(c);
            }
            else if (c == '*') {
                _state = OP_MUL;
                end = _location;
                _location.moveToNext(c);
                list_token.add(c);
            }
            else if (c == '/') {
                _state = OP_DIV;
                end = _location;
                _location.moveToNext(c);
                list_token.add(c);
            }
            else if (c == '|') {
                _state = OP_OR;
                end = _location;
                _location.moveToNext(c);
                list_token.add(c);
            }
            else if (c == '&') {
                _state = OP_AND;
                end = _location;
                _location.moveToNext(c);
                list_token.add(c);
            }
            else if (c == '<') {
                _state = OP_LT;
                end = _location;
                _location.moveToNext(c);
                list_token.add(c);
            }
            else if (c == '>') {
                _state = OP_GT;
                end = _location;
                _location.moveToNext(c);
                list_token.add(c);
            }
            else if (c == '^') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::BIT_XOR, "^", 1, SourceRange(begin, end));
            }
            else if (c == ';') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::ENDL, "", 0, SourceRange(begin, end));
            }
            else if (c == '\0') {
                _state = INIT;
                end = _location;
                _location.moveToNext(c);
                return new Token(TokenType::ENDMARKER, "", 0, SourceRange(begin, end));
            }
            else if (c == '"') {
                list_token.add(c);
                _state = STRING;
                end = _location;
                _location.moveToNext(c);
            }
            else if (c == '\'') {
                _state = CHAR;
                end = _location;
                _location.moveToNext(c);
            }
            else if (c == ':') {
                list_token.add(c);
                _state = OP_COLON;
                end = _location;
                _location.moveToNext(c);
            }
        }
        else if (_state == INTEGER) {
            if (c >= '0' && c <= '9') {
                _stream->advance();
                list_token.add(c);
                end = _location;
                _location.moveToNext(c);
            }
            else {
                return make_new_token(TokenType::INT, list_token, begin, end, '\0');
            }
        }
        else if (_state == NAME) {
            if ((c >= 'a' && c <= 'z') 
                || (c >= 'A' && c <= 'Z') 
                || c == '_'
                || (c >= '0' && c <= '9')) {
                    _stream->advance();
                    list_token.add(c);
                    end = _location;
                    _location.moveToNext(c);
            }
            else {
                return make_new_token(TokenType::NAME, list_token, begin, end, '\0');
            }
        }
        else if (_state == OP_EQ) {
            if (c == '>') {
                _stream->advance();
                end = _location;
                _location.moveToNext(c);
                return make_new_token(TokenType::LAM_DEF, list_token, begin, end, c);
            }
            else if (c == '=') {
                _stream->advance();
                end = _location;
                _location.moveToNext(c);
                return make_new_token(TokenType::EQUAL, list_token, begin, end, c);
            }
            else {
                return make_new_token(TokenType::ASN, list_token, begin, end, '\0');
            }
        }
        else if (_state == OP_SLASH) {
            // TODO
        }
        else if (_state == OP_ADD) {
            if (c == '=') {
            }
            else {
                return make_new_token(TokenType::PLUS, list_token, begin, end, '\0');
            }

        }
        else if (_state == OP_SUB) {
            if (c == '=') {
            }
            else if (c == '>') {
                _stream->advance();
                end = _location;
                _location.moveToNext(c);
                return make_new_token(TokenType::ARROW, list_token, begin, end, c);
            }
            else {
                return make_new_token(TokenType::MINUS, list_token, begin, end, '\0');
            }
        }
        else if (_state == OP_MUL) {
            if (c == '=') {
            }
            else {
                return make_new_token(TokenType::MULT, list_token, begin, end, '\0');
            }
        }
        else if (_state == OP_DIV) {
            if (c == '=') {
            }
            else {
                return make_new_token(TokenType::DIV, list_token, begin, end, '\0');
            }
        }
        else if (_state == OP_OR) {
            if (c == '|') { // logic or
                _stream->advance();
                end = _location;
                _location.moveToNext(c);
                return make_new_token(TokenType::LOGIC_OR, list_token, begin, end, c);
            }
            else {
                return make_new_token(TokenType::BIT_OR, list_token, begin, end, '\0');
            }
        }
        else if (_state == OP_AND) {
            if (c == '&') { // logic and
                _stream->advance();
                end = _location;
                _location.moveToNext(c);
                return make_new_token(TokenType::LOGIC_AND, list_token, begin, end, c);
            }
            else {
                return make_new_token(TokenType::BIT_AND, list_token, begin, end, '\0');
            }
        }
        else if (_state == OP_LT) {
            if (c == '<') {
                _stream->advance();
                end = _location;
                _location.moveToNext(c);
                return make_new_token(TokenType::LEFT_SHIFT, list_token, begin, end, c);
            }
            else {
                return make_new_token(TokenType::LT, list_token, begin, end, '\0');
            }
        }
        else if (_state == OP_GT) {
            if (c == '>') {
                _stream->advance();
                end = _location;
                _location.moveToNext(c);
                return make_new_token(TokenType::RIGHT_SHIFT, list_token, begin, end, c);
            }
            else {
                return make_new_token(TokenType::GT, list_token, begin, end, '\0');
            }
        }
        else if (_state == STRING) {
            _stream->advance();
            list_token.add(c);
            end = _location;
            _location.moveToNext(c);
            if (c == '"') {
                return make_new_token(TokenType::STRING, list_token, begin, end, '\0');
            }
        }
        else if (_state == CHAR) {
            _stream->advance();
            end = _location;
            _location.moveToNext(c);
            if (c != '\'') {
                list_token.add(c);
            }
            else {
                return make_new_token(TokenType::CHAR, list_token, begin, end, '\0');
            }
        }
        else if (_state == OP_COLON) {
            if (c == ':') {
                _stream->advance();
                end = _location;
                _location.moveToNext(c);
                return make_new_token(TokenType::QUALIFIER, list_token, begin, end, c);
            }
            else {
                return make_new_token(TokenType::COLON, list_token, begin, end, '\0');
            }
        }
    }

    return NULL;
}

Token* Lexer::make_new_token(TokenType tt, ArrayList<char>& alist, SourceLocation begin, SourceLocation end, char c) {
    if (c != 0) {
        alist.add(c);
    }

    Token* t = new Token(tt, alist.value(), alist.size(), SourceRange(begin, end));
    if (tt == TokenType::CHAR) {
        char c = filter_char(alist.value(), alist.size());
        *t->_value = c;
    }
    alist.clear();
    _state = INIT;

    if (tt == TokenType::NAME) {
        return filter_keyword(t);
    }
    return t;
}

void Token::print() {
    printf("%s: ", toString(_tt));
    for (int i = 0; i < _length; i++) {
        printf("%c", _value[i]);
    }
    printf("\n");
}

int Token::cmp(const char* s) {
    if (_length != (int)strlen(s)) {
        return -1;
    }
    
    for (int i = 0; i < _length; i++) {
        if (_value[i] != *(s+i)) {
            return -1;
        }
    }
    
    return 0;
}

